import os.path

from transformers import AutoTokenizer
from datasets import load_from_disk, load_dataset
import json


# HF_ENDPOINT=https://hf-mirror.com
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased", cache_dir="./MNIST/cache", force_download=False)
print(tokenizer)
print(tokenizer("What is your name?", "My name is William"))


def prepare_train_features(all_questions, all_contexts, all_answers):
    # Some of the questions have lots of whitespace on the left, which is not useful and will make the
    # truncation of the context fail (the tokenized question will take a lots of space). So we remove that
    # left whitespace
    all_questions = [q.lstrip() for q in all_questions]

    # Tokenize our examples with truncation and padding, but keep the overflows using a stride. This results
    # in one example possible giving several features when a context is long, each of those features having a
    # context that overlaps a bit the context of the previous feature.
    tokenized_examples = tokenizer(
        all_questions,
        all_contexts,
        truncation='only_second',
        max_length=384,
        stride=128,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding='max_length',
    )

    # Since one example might give us several features if it has a long context, we need a map from a feature to
    # its corresponding example. This key gives us just that.
    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
    # The offset mappings will give us a map from token to character position in the original context. This will
    # help us compute the start_positions and end_positions.
    offset_mapping = tokenized_examples.pop("offset_mapping")

    # Let's label those examples!
    tokenized_examples["start_positions"] = []
    tokenized_examples["end_positions"] = []

    for i, offsets in enumerate(offset_mapping):
        # We will label impossible answers with the index of the CLS token.
        input_ids = tokenized_examples["input_ids"][i]
        cls_index = input_ids.index(tokenizer.cls_token_id)

        # Grab the sequence corresponding to that example (to know what is the context and what is the question).
        sequence_ids = tokenized_examples.sequence_ids(i)

        # One example can give several spans, this is the index of the example containing this span of text.
        sample_index = sample_mapping[i]
        answers = all_answers[sample_index]
        # If no answers are given, set the cls_index as answer.
        if answers is None:
            tokenized_examples["start_positions"].append(cls_index)
            tokenized_examples["end_positions"].append(cls_index)
        else:
            # Start/end character index of the answer in the text.
            start_char = answers["answer_start"]
            end_char = start_char + len(answers["text"])

            # Start token index of the current span in the text.
            token_start_index = 0
            while sequence_ids[token_start_index] != 1:
                token_start_index += 1

            # End token index of the current span in the text.
            token_end_index = len(input_ids) - 1
            while sequence_ids[token_end_index] != 1:
                token_end_index -= 1

            # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).
            if not (offsets[token_start_index][0] <= start_char
                    and offsets[token_end_index][1] >= end_char):
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
            else:
                # Otherwise move the token_start_index and token_end_index to the two ends of the answer.
                # Note: we could go after the last offset if the answer is the last word (edge case).
                while token_start_index < len(offsets) and offsets[
                        token_start_index][0] <= start_char:
                    token_start_index += 1
                tokenized_examples["start_positions"].append(
                    token_start_index - 1)
                while offsets[token_end_index][1] >= end_char:
                    token_end_index -= 1
                tokenized_examples["end_positions"].append(token_end_index + 1)

    return tokenized_examples


# dataset = load_dataset("lhoestq/squad", cache_dir="./MNIST/cache")
all_contexts = []
all_questions = []
all_answers = []
with open("./MNIST/squad/dev-v2.0.json", "r", encoding="utf-8") as j:
    json = json.load(j)
    print(json["data"][0]["paragraphs"][0])
    for article in json["data"]:
        for paragraph in article["paragraphs"]:
            context = paragraph["context"]
            for qa in paragraph["qas"]:
                question = qa["question"]
                answers = qa["answers"]
                all_contexts.append(context)
                all_questions.append(question)

                if len(answers) > 1:
                    all_answers.append(answers[0])
                else:
                    all_answers.append(None)
examples = prepare_train_features(all_questions, all_contexts, all_answers)
print(f"examples:{examples}")

for i in range(9):
    input_ids = examples["input_ids"][i]
    start_positions = examples["start_positions"][i]
    end_positions = examples["end_positions"][i]
    print(f"问题和文本：{tokenizer.decode(input_ids)}")
    print(f"答案：{start_positions} {end_positions} {tokenizer.decode(input_ids[start_positions:end_positions])}")
    print(f"原始答案：{all_answers[i]}")



